open Types
open Sast
open Printer
open Printf
open Bigarray
open Eval
open Dimlist

exception Exception of string
exception UnresolvedParam

type resolve_env = {
    unres_symbols: int list;
}

let create_env (p:Sast.prog) =
      let unres = IntMap.fold
          (fun key param li -> 
             match param with
               { pa_type=MatrixUnknownSize; pa_num=pa_num} -> pa_num :: li
             | _ ->  li )
          p.pparams [] 
      in
      let unres = IntMap.fold
          (fun gsid n li -> 
             match n with
               { t=MatrixUnknownSize; } -> gsid :: li
             | _ ->  li )
          p.pneurons unres 
      in
      { unres_symbols=unres }

(* Returns the number of symbols for which the sizes are unknown *)
let ambiguity { unres_symbols=unres_symbols } { syn=syn } =
    List.length unres_symbols + 
    (IntMap.fold
       (fun key syn num -> match syn.sexpr with _,MatrixUnknownSize->num+1 | _->num ) syn 0)




(* Inside index *)
let rec retrieve_forvar = function
    Id(varn) -> [varn]
  | Intgr(_) -> []
  | SizeOf(_,_) -> [] (* for-variables are not allowed in sizeof *)
  | Index((e1,t1),_) ->
      (retrieve_forvar e1)
  | Binop((e1,t1),op,(e2,t2)) ->
      let varn1 = retrieve_forvar e1 in
      let varn2 = retrieve_forvar e2 in 
        List.append varn1 varn2 
  | Unop(op,(e1,t1)) ->
      (retrieve_forvar e1)
  | Span(e1,e2,e3) -> [] (* for-variables not allowed in Span *)
  | NeuronRef(_,_,_) -> assert false
  | Param(_,_) -> assert false
  | ExtNeuron(_,_,_,_,_) -> assert false
  | ForExpr(_,_) -> assert false
  | Conv(_,_) -> assert false
  | Filter(_,_,_) -> assert false
  | WghtRef(_,_,_) -> assert false
  | ActRef(_,_,_,_) -> assert false
  | Float(_) -> assert false

(* shortens the dimension lists when possible, ex. [256,256,1] becomes
 * [256,256], however, must have atleast 1 element, ex. [1,1] becomes [1]
 *)           
let shorten_diml diml = 
  match 
     (fst 
        (List.fold_right 
           (fun d (li,do_reduce) ->
              if do_reduce && d = 1 then
                (li,true)
              else
                (d :: li,false)) 
           diml
           ([],true)
        )
     )
  with [] -> [1]
    | diml -> diml
      
type opt_int = Int of int | NoInt

let rec resolve_looped_dim dexpr sfor =
  let max_int nl = 
    List.fold_left 
      (fun (x:int) (y:int) -> if (x > y) then x else y) (-1) nl
  in 
  let varnl = retrieve_forvar dexpr in
  match (dexpr, varnl, sfor) with
    | _,[varn], ForExpr(id,span)::forlist -> 
         ( if varn <> id then 
             resolve_looped_dim dexpr forlist 
           else
           (
             let (nums : int list) = Eval.nums_of_span span in 
               max_int nums
           )
         )
    | Span(_,_,_),[],_ -> 
        let (nums : int list) = Eval.nums_of_span dexpr in 
          max_int nums
    | _,_,ForExpr(id,span)::_ -> 
        raise( Exception 
                 ("Synapse does not currently support having multiple "^ 
                 "for-variables in a single index dimension.")) 
    | _,_,sast::_ -> 
        raise( Exception ("Error: Invalid expression type "^(string_of_exprd sast)^" in for-list!"))

(** Calculates the dimension of the synaptic targets.  If no index is
  * specified for the destination, then it is just set to expr_diml.  Otherwise, the dimension
  * is calculated from the destination index when the for-variables are specified.
 *)
let resolve_synap_dest sfor indl expr_diml = 
  let diml' = List.map (fun d -> resolve_looped_dim d sfor) indl in
    match diml' with 
        d::dl -> (shorten_diml diml')
      | _ -> expr_diml

(* returns modified program and list of unresolved parameters *)
let resolve_symbol_ref p pnum =
   let synap = IntMap.find pnum p.syn in 
     match synap with 
         { sexpr=(sexprd,Types.MatrixUnknownSize); sfor=sfor} ->
           p,Int(pnum)      (* can't do anything, but push pnum as being unresolved *)
       | { sexpr=sexprd,Types.Int; sfor=sfor} ->
           assert false (* compiler error *)

       | { sdest=Param(pnum,indl); sexpr=sexprd,Types.Matrix(expr_diml); sfor=sfor} ->
           (
           print_endline (Printer.string_of_synap pnum synap); 
           assert ( IntMap.mem pnum p.pparams = true ); 
           let pam = IntMap.find pnum p.pparams in
             match pam with
                 { pa_type=Matrix(li) } -> p,NoInt (* nothing to do *)
               | { pa_type=MatrixUnknownSize } ->
                   (
                   try
                   let diml' = resolve_synap_dest sfor indl expr_diml in
                    Printf.printf "Expr_diml=[%s]\n" (Printer.string_of_int_list expr_diml ","); 
                    Printf.printf "New diml=[%s]\n" (Printer.string_of_int_list diml' ","); 
                   let params = IntMap.add pnum 
                                  { pa_num=pnum; pa_type=Types.Matrix( diml' ); 
                                    pa_ptype=pam.pa_ptype; pa_io=pam.pa_io; pa_data=Unallocated }
                                  p.pparams
                   in
                     { syn=p.syn; pparams=params; 
                       actfuns=p.actfuns; pneurons=p.pneurons; kerfuns=p.kerfuns;
                     }, NoInt
                   with UnresolvedParam -> p, Int(pnum)
                   )
               | _ ->
                     raise ( Exception "For-loop not yet completed.")
           )

       | { sdest=NeuronRef(gsid,name,indl); sexpr=sexprd,Types.Matrix(expr_diml); sfor=sfor} ->
           (
           print_endline (Printer.string_of_synap pnum synap); 
           assert ( IntMap.mem pnum p.pneurons = true ); 
           let n = IntMap.find pnum p.pneurons in
             match n with
                 { t=Matrix(li) } -> p,NoInt (* nothing to do *)
               | { t=MatrixUnknownSize } ->
                   (
                   try
                   let diml' = resolve_synap_dest sfor indl expr_diml in
                   let neurons' = IntMap.add gsid 
                                  { t=Types.Matrix( diml' ); 
                                    v=n.v;
                                    n_name = n.n_name; n_mname=n.n_mname }
                                  p.pneurons
                   in
                     { syn=p.syn; pparams=p.pparams; 
                       actfuns=p.actfuns; kerfuns=p.kerfuns; pneurons=neurons' }, NoInt
                   with UnresolvedParam -> p, Int(pnum)
                   )
               | _ ->
                     raise ( Exception "For-loop not yet completed.")
           )
       | _ -> raise ( Exception "Parameter misassociated with external neuron reference.")

(* calculate the "span" of the dimension given the restriction of the index 
 *  TODO: Need to call something other than Eval.nums_of_span to handle the case
 *  when an for-index is in the span.
 *)
let index_span exprd = match exprd with 
    Sast.Span(e1,e2,e3) -> List.length (Eval.nums_of_span exprd)
  | _ -> 1

let rec eval_expr_diml indl diml p = 
  match (indl, diml) with
      [],[] -> []
    | (ind :: indl'),(dim :: diml') ->
        (index_span ind) :: (eval_expr_diml indl' diml' p)
    | [],diml -> (shorten_diml diml)
    | _ ->
        raise ( Exception "Index specifies more dimensions than what is declared.")

let rec spans_of_indl indl p = 
  match indl with
      [] -> []
    | (ind :: indl') -> (index_span ind) :: (spans_of_indl indl' p)
                
(* returns the hopefully less ambiguious expression and the number
* of expressions that have had their types resolves *)
let rec resolve_expr expr p = match (expr:Sast.expr) with
    _,Types.Int -> expr,0
  | _,Types.Matrix(_) -> expr,0
  | Sast.Unop(op,e1), Types.MatrixUnknownSize ->
      let (e1',t1),amb1 = resolve_expr e1 p in
      let t',a = match t1 with
          Types.Int -> Types.Matrix([1]), 1
        | Types.Matrix(diml) -> Types.Matrix(diml),1
        | Types.MatrixUnknownSize -> Types.MatrixUnknownSize,0
      in
        (Sast.Unop(op,(e1',t1)), t'), (amb1+a)

  | Sast.Binop(e1,op,e2), Types.MatrixUnknownSize ->
      let (e1',t1),amb1 = resolve_expr e1 p in
      let (e2',t2),amb2 = resolve_expr e2 p in
      let t',a = match (t1,t2) with
          (Types.MatrixUnknownSize,_) -> Types.MatrixUnknownSize, 0 (* 0=amb resolution at current node *)
        | (_,Types.MatrixUnknownSize) -> Types.MatrixUnknownSize, 0
        | (Types.Matrix(diml),Types.Int) -> Types.Matrix(diml), 0
        | (Types.Int, Types.Matrix(diml)) -> Types.Matrix(diml), 0
        | (Types.Matrix(diml),Types.Matrix([1])) -> Types.Matrix(diml), 1
        | (Types.Matrix([1]),Types.Matrix(diml)) -> Types.Matrix(diml), 1
        | (t1,t2) -> 
            if t1=t2 then t1,1
            else 
              raise (Exception 
                       (Printf.sprintf 
                          "Binop has non-compatible types %s and %s." 
                          (Printer.string_of_type t1) 
                          (Printer.string_of_type t2)) )
      in 
        print_endline (Printer.string_of_type t1); 
        print_endline (Printer.string_of_type t2);
        (Sast.Binop((e1',t1),op,(e2',t2)), t'), (amb1+amb2+a)
  | Sast.Conv(e1,e2), Types.MatrixUnknownSize ->
      let (e1',t1),amb1 = resolve_expr e1 p in
      let (e2',t2),amb2 = resolve_expr e2 p in
      let t,t2',a = match (t1,t2) with
          (Types.MatrixUnknownSize,t) -> Types.MatrixUnknownSize, t, 0
        | (Types.Matrix(diml1),Types.MatrixUnknownSize) -> Types.Matrix([1]),Types.Matrix(diml1),2
        | (Types.Matrix(diml1),Types.Matrix(diml2)) ->
            assert (diml1=diml2);
            Types.Matrix([1]),Types.Matrix(diml1),1
        | (Types.Int, MatrixUnknownSize) -> Types.Matrix([1]),Types.Matrix([1]),2
        | (Types.Int, Matrix([1])) -> Types.Matrix([1]),Types.Matrix([1]),1
        | (Types.Int, _ ) -> assert false
        | ( _ ,Types.Int) -> assert false
      in
        (Sast.Conv((e1',t1),(e2',t2')), t), (amb1+amb2+a)

  | Sast.Filter(e1,e2,incr), Types.MatrixUnknownSize ->
      let (e1',t1),amb1 = resolve_expr e1 p in
      let (e2',t2),amb2 = resolve_expr e2 p in
      let t,a = 
        match (t1,t2) with
          (Types.MatrixUnknownSize,_) -> Types.MatrixUnknownSize,0
        | (_,Types.MatrixUnknownSize) -> Types.MatrixUnknownSize,0
        | (Types.Matrix(diml1),Types.Matrix(diml2)) ->
            let diml = 
              Dimlist.map2 
                (fun d1 d2 -> 
                   let ret = d1 - d2 + 1
                   in
                     if ret > 0 then ret
                     else
                       raise (Exception "The first parameter of filter must be larger than the second parameter in all dimensions.")
                )
                diml1 diml2
            in Types.Matrix(diml), 1
        | _ -> raise (Exception "The built in filter function can only take matrices parameters.")
      in
      let rexpr = (Sast.Filter((e1',t1),(e2',t2),incr), t) 
      in
        Printf.printf "Resolved filter to be of type %s.\n" (Printer.string_of_type t); 
        rexpr, (amb1+amb2+a)

  | Sast.WghtRef(gsid,name,fparams), Types.MatrixUnknownSize -> (* resolution occurs by Conv not here*)
      (Sast.WghtRef(gsid,name,fparams),Types.MatrixUnknownSize), 0
  | Sast.ActRef(gsid,name,e1,fparams), Types.MatrixUnknownSize ->
      let (e1',t1),amb1 = resolve_expr e1 p in
      let t,a =
        (
        match t1 with
            Types.MatrixUnknownSize -> t1,0
          | Types.Matrix([1]) -> t1,1
          | Types.Matrix(diml) -> 
              raise ( Exception "A non-scalar expression cannot be passed to an activation function.")
          | Types.Int -> assert false
        )
      in
        (Sast.ActRef(gsid,name,(e1',t1),fparams),t),(amb1+a)

  | Sast.Param(pnum,indl), Types.MatrixUnknownSize ->
      let pam = IntMap.find pnum p.pparams in
        (Printf.printf 
           "Trying to resolve ambiguity of $%d reference with %s.\n" 
           pnum (string_of_t pam.pa_type));
        (
        match pam.pa_type with
            Types.MatrixUnknownSize-> (Sast.Param(pnum,indl), MatrixUnknownSize), 0
          | Types.Matrix(diml) -> 
              let expr_diml = shorten_diml (eval_expr_diml indl diml p) in 
                (Param(pnum,indl), Matrix(expr_diml)), 1
          | Types.Int -> assert false
        )
  | Sast.ExtNeuron(gsid, modref, [], nname, nindl), Types.MatrixUnknownSize ->
      let {t=nt} = IntMap.find gsid p.pneurons in
        (
        match nt with
            Types.Int -> assert false
          | Types.MatrixUnknownSize -> 
              ( Sast.ExtNeuron(gsid, modref, [], nname, nindl), MatrixUnknownSize), 0
          | Types.Matrix(diml) ->
              let expr_diml = shorten_diml (eval_expr_diml nindl diml p) in 
              ( Sast.ExtNeuron(gsid, modref, [], nname, nindl), Matrix(expr_diml) ), 1
        )
  | Sast.Index(e1,indl), Types.MatrixUnknownSize ->
      let (e1',t1),amb1 = resolve_expr e1 p in
      let spanl = spans_of_indl indl p in
        Printf.printf "Resolving index to matrix of size [";
        List.iter 
          (Printf.printf " %d") spanl;
        Printf.printf "].\n";
        (Sast.Index( (e1',t1), indl ), Types.Matrix(spanl)),1

  | Sast.NeuronRef(gsid, nname, nindl), Types.MatrixUnknownSize ->
      let {t=nt} = IntMap.find gsid p.pneurons in
        (
        match nt with
            Types.Int -> assert false
          | Types.MatrixUnknownSize -> 
              ( Sast.NeuronRef(gsid, nname, nindl), MatrixUnknownSize), 0
          | Types.Matrix(diml) ->
              let expr_diml = shorten_diml (eval_expr_diml nindl diml p) in 
              ( Sast.NeuronRef(gsid, nname, nindl), Matrix(expr_diml) ), 1
        )
  | e -> raise (Exception ("No resolution for " ^ (string_of_expr e)))

(* Returns synapse with expression with possibly less ambiguity and
* the number of expressions in the SAST that have been resolved. *)
let resolve_syn s p = 
  let sexpr',ambigreduction = resolve_expr s.sexpr p
  in 
    { sdest=s.sdest; sexpr=sexpr'; sfor=s.sfor }, ambigreduction


let rec resolve_sizes p renv =
  Printf.printf "Ambiguity = %d\n" (ambiguity renv p);
   if ambiguity renv p = 0 then p
   else
    let unr_syms, p = List.fold_left 
                       (fun (li,p) pnum -> match (resolve_symbol_ref p pnum) with
                             p,NoInt     -> li, p
                           | p,Int(pnum) -> pnum::li,p )
                       ([],p) renv.unres_symbols
    in  
      print_endline "After resolve_symbol_ref:"; 
      print_endline (IntMap.fold (fun k pam str -> (string_of_param k pam) ^ str ) p.pparams "");
    let synaps',ambigreduction = 
      IntMap.fold 
        (fun k s (syns,ar) -> 
                     let syn,ared = resolve_syn s p in
                       IntMap.add k syn syns,(ar+ared)) 
        p.syn (IntMap.empty,0)
    in
     print_endline "After calling resolve_syn:";
     print_endline 
       (IntMap.fold 
          (fun k syn str -> (string_of_synap k syn) ^ str ) synaps' "");
    let p = { syn=synaps'; pparams=p.pparams; 
              actfuns=p.actfuns; kerfuns=p.kerfuns; pneurons=p.pneurons } 
    in 
    let renv' = { unres_symbols = unr_syms }
    in
       if ambiguity renv p <= ambiguity renv' p && ambigreduction = 0 then
       (
          print_string (Printer.string_of_program p);
          raise ( Exception "Unable to resolve ambiguity of symbol dimensions." );
       )
       else
          resolve_sizes p renv'

let allocate_neuron_array gsid p n = match n with
    { t=Matrix(diml); v=v } -> 
      let reqd_size = List.fold_left (fun a dim -> a*dim) 1 diml in
      let v' = Array.make_matrix reqd_size 2 0.0 in
        { t = Matrix(diml); v = v'; n_name=n.n_name; n_mname=n.n_mname }
  | { t=Types.Int } -> assert false;
  | { t=Types.MatrixUnknownSize } -> 
      if IntMap.mem gsid p.syn = false then 
        raise ( Exception (Printf.sprintf "Neuron %d is not defined by a synapse." gsid));
      let { sexpr=(expr,t) } = IntMap.find gsid p.syn in
        match t with
            Types.Matrix(diml) -> 
              let reqd_size = List.fold_left (fun a dim -> a*dim) 1 diml in 
              let v' = Array.make_matrix reqd_size 2 0.0 in 
                { t = Types.Matrix(diml); v = v'; 
                  n_name=n.n_name; n_mname=n.n_mname }
          | _ -> assert false (* compiler error, all synapses have been resolved *)

let create_image_data d1 d2 d3 =  (* should allow different alignments *)
  let data' = Array3.create 
                Bigarray.int8_unsigned 
                Bigarray.c_layout d1 d2 d3 (* order of data is y-axis x-axis nChannels *)
  in
    Array3.fill data' 0;
    data'

let allocate_param_data pam = match pam with
    { pa_io=SpaceDelimFile(_,_); pa_type=Matrix(diml) } -> 
      let data' =
        let totaldim = 
          List.fold_left (fun a d -> a*d) 1 diml
        in
          Array.make totaldim 0.0
      in
      { pa_num=pam.pa_num; pa_type=pam.pa_type; 
        pa_ptype=pam.pa_ptype; pa_io=pam.pa_io;
        pa_data=FloatArray(data') }

  | { pa_type=Matrix([d1; d2; d3]); pa_io=OutPpmSeq(fpat) } ->
        { pa_num=pam.pa_num; pa_type=pam.pa_type;
          pa_ptype=pam.pa_ptype;
          pa_io=OutPpmSeq(fpat); pa_data=ImageC3(create_image_data d1 d2 d3) }

  | { pa_type=Matrix([d1; d2; d3]); pa_io=InPpm(fpat,nreads) } ->
        { pa_num=pam.pa_num; pa_type=pam.pa_type;
          pa_ptype=pam.pa_ptype;
          pa_io=InPpm(fpat,nreads);
          pa_data=ImageC3(create_image_data d1 d2 d3)
        }

  | { pa_type=Types.MatrixUnknownSize; pa_num=pnum } 
    -> raise( Exception (Printf.sprintf "Param $%d has unresolved size." pnum))
  | { pa_type=Types.Int } -> raise( Exception "Param cannot have type int.")
  | { pa_type=Matrix(diml) } ->
      raise( Exception 
               (Printf.sprintf 
                  "Parameter must have exactly 3 dimensions, but has dimensionality of %s.\n" 
                  (Printer.string_of_type pam.pa_type))
      )


let program p =
   let p = resolve_sizes p (create_env p) in
     print_endline "During Translation2:";
     print_endline (Printer.string_of_program p);
   let neurons' = IntMap.mapi (fun key n -> allocate_neuron_array key p n) p.pneurons in
   let params' = IntMap.map allocate_param_data p.pparams in
     { syn=p.syn; pneurons=neurons'; pparams=params'; actfuns=p.actfuns; kerfuns=p.kerfuns }
